package com.example.mqtt.client;

import io.netty.buffer.ByteBuf;
import io.netty.channel.EventLoop;
import io.netty.handler.codec.mqtt.*;
import io.netty.util.concurrent.ScheduledFuture;
import lombok.Data;

import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;
import java.util.function.BiConsumer;
import java.util.function.Consumer;

public final class MqttPendingObject {

    static final class RetransmissionHandler<T extends MqttMessage> {

        private ScheduledFuture<?> timer;
        private int timeout = 10;
        private BiConsumer<MqttFixedHeader, T> handler;
        private T originalMessage;

        void start(EventLoop eventLoop) {
            if(this.handler == null){
                throw new NullPointerException("handler");
            }
            this.timeout = 10;
            this.startTimer(eventLoop);
        }

        private void startTimer(EventLoop eventLoop){
            this.timer = eventLoop.schedule(() -> {
                this.timeout += 5;
                MqttFixedHeader fixedHeader = new MqttFixedHeader(this.originalMessage.fixedHeader().messageType(), true, this.originalMessage.fixedHeader().qosLevel(), this.originalMessage.fixedHeader().isRetain(), this.originalMessage.fixedHeader().remainingLength());
                handler.accept(fixedHeader, originalMessage);
                startTimer(eventLoop);
            }, timeout, TimeUnit.SECONDS);
        }

        void stop(){
            if(this.timer != null){
                this.timer.cancel(true);
            }
        }

        void setHandle(BiConsumer<MqttFixedHeader, T> runnable) {
            this.handler = runnable;
        }

        void setOriginalMessage(T originalMessage) {
            this.originalMessage = originalMessage;
        }
    }

    final static class MqttIncomingQos2Publish {
        private final MqttPublishMessage incomingPublish;
        private final RetransmissionHandler<MqttMessage> retransmissionHandler = new RetransmissionHandler<>();

        MqttIncomingQos2Publish(MqttPublishMessage incomingPublish, MqttMessage originalMessage) {
            this.incomingPublish = incomingPublish;
            this.retransmissionHandler.setOriginalMessage(originalMessage);
        }

        MqttPublishMessage getIncomingPublish() {
            return incomingPublish;
        }

        void startRetryPubRec(EventLoop eventLoop, Consumer<Object> sendPacket) {
            this.retransmissionHandler.setHandle((fixedHeader, originalMessage) ->
                    sendPacket.accept(new MqttMessage(fixedHeader, originalMessage.variableHeader())));
            this.retransmissionHandler.start(eventLoop);
        }
        void stopRetryPubRec() {
            this.retransmissionHandler.stop();
            incomingPublish.release();
        }
    }


    @Data
    static final class MqttPendingPublish {

        private final int messageId;
        private final CompletableFuture<Void> future;
        private final ByteBuf payload;
        private final MqttPublishMessage message;
        private final MqttQoS qos;

        private final RetransmissionHandler<MqttPublishMessage> publishRetransmissionHandler = new RetransmissionHandler<>();
        private final RetransmissionHandler<MqttMessage> pubrelRetransmissionHandler = new RetransmissionHandler<>();

        private boolean sent = false;

        MqttPendingPublish(int messageId, CompletableFuture<Void> future, ByteBuf payload, MqttPublishMessage message, MqttQoS qos) {
            this.messageId = messageId;
            this.future = future;
            this.payload = payload;
            this.message = message;
            this.qos = qos;

            this.publishRetransmissionHandler.setOriginalMessage(message);
        }

        void startRetryPublish(EventLoop eventLoop, Consumer<Object> sendPacket) {
            this.publishRetransmissionHandler.setHandle(((fixedHeader, originalMessage) ->
                    sendPacket.accept(new MqttPublishMessage(fixedHeader, originalMessage.variableHeader(), this.payload.retain()))));
            this.publishRetransmissionHandler.start(eventLoop);
        }

        void stopRetryPublish() {
            this.publishRetransmissionHandler.stop();
            payload.release();
        }

        void setPubrelMessage(MqttMessage pubrelMessage) {
            this.pubrelRetransmissionHandler.setOriginalMessage(pubrelMessage);
        }

        void startRetryPubrel(EventLoop eventLoop, Consumer<Object> sendPacket) {
            this.pubrelRetransmissionHandler.setHandle((fixedHeader, originalMessage) ->
                    sendPacket.accept(new MqttMessage(fixedHeader, originalMessage.variableHeader())));
            this.pubrelRetransmissionHandler.start(eventLoop);
        }

        void stopRetryPubrel() {
            this.pubrelRetransmissionHandler.stop();
        }
    }


    @Data
    static final class MqttPendingSubscription {

        private final CompletableFuture<Void> future;
        private final String topic;
        private final MqttSubscribeMessage subscribeMessage;

        private final RetransmissionHandler<MqttSubscribeMessage> retransmissionHandler = new RetransmissionHandler<>();

        private boolean sent = false;

        MqttPendingSubscription(CompletableFuture<Void> future, String topic, MqttSubscribeMessage message) {
            this.future = future;
            this.topic = topic;
            this.subscribeMessage = message;
            this.retransmissionHandler.setOriginalMessage(message);
        }

        void startRetrySubscribe(EventLoop eventLoop,Consumer<Object> sendPacket) {
            if(this.sent){ //If the packet is sent, we can start the retransmit timer
                this.retransmissionHandler.setHandle((fixedHeader, originalMessage) ->
                        sendPacket.accept(new MqttSubscribeMessage(fixedHeader, originalMessage.variableHeader(), originalMessage.payload())));
                this.retransmissionHandler.start(eventLoop);
            }
        }

        void stopRetrySubscribe(){
            this.retransmissionHandler.stop();
        }
    }

    static final class MqttPendingUnsubscription {

        private final CompletableFuture<Void> future;
        private final String topic;

        private final RetransmissionHandler<MqttUnsubscribeMessage> retransmissionHandler = new RetransmissionHandler<>();

        MqttPendingUnsubscription(CompletableFuture<Void> future, String topic, MqttUnsubscribeMessage unsubscribeMessage) {
            this.future = future;
            this.topic = topic;

            this.retransmissionHandler.setOriginalMessage(unsubscribeMessage);
        }

        CompletableFuture<Void> getFuture() {
            return future;
        }

        String getTopic() {
            return topic;
        }

        void startRetryUnSub(EventLoop eventLoop, Consumer<Object> sendPacket) {
            this.retransmissionHandler.setHandle((fixedHeader, originalMessage) ->
                    sendPacket.accept(new MqttUnsubscribeMessage(fixedHeader, originalMessage.variableHeader(), originalMessage.payload())));
            this.retransmissionHandler.start(eventLoop);
        }

        void stopRetryUnSub(){
            this.retransmissionHandler.stop();
        }
    }

}
